package chen.bupt.datastructure.kdtree

import java.io.File

import chen.bupt.help.kdtree.Rect

import scala.collection.mutable
import scala.collection.mutable.ListBuffer
import scala.io.Source


/**
 * Created by chenlingpeng on 2014/11/19.
 */
class KdTree(k: Int) {
  var root: Option[KdNode] = None

  def init(nodes: List[Node]) = {
    import chen.bupt.datastructure.kdtree.KdTree._
    root = createKdTree(nodes, 0, k)
  }

  def addNode(node: Node) = {
    // TODO: 为用户训练设计的接口
  }



  var count = 0
  private def searchNode(value: Vector[Double], node: KdNode, results: mutable.TreeSet[KdTmpNode], visited: mutable.HashSet[KdNode], K: Int): Unit = {
//    var start = System.currentTimeMillis()
//    if(count%100==0)
//      println(count)
    count+=1
    visited.add(node)

    val nodeDist = Node.euclideanDistance(value, node.node.vec)
    val tmpNode = KdTmpNode(node.node.id, nodeDist)
    if (results.size == K) {
      val lastNode: KdTmpNode = results.last
      val lastDist: Double = lastNode.dist
      if (lastDist > nodeDist) {
        results.remove(lastNode)
        results.add(tmpNode)
      }
    } else {
      results.add(tmpNode)
    }
//    if(results.isEmpty || results.get.dist>nodeDist){
//      results.set(tmpNode)
//    }

    val lastNode: KdTmpNode = results.last
    val lastDist: Double = lastNode.dist
//    println(System.currentTimeMillis()-start)
//    start = System.currentTimeMillis()

    if(node.leftNode.isDefined && !visited.contains(node.leftNode.get)){
      visited.add(node.leftNode.get)
      // - 逼近下方
//      println(System.currentTimeMillis()-start)
      if(value(node.depth%node.k)-lastDist<node.node.vec(node.depth%node.k)){
        searchNode(value, node.leftNode.get, results, visited, K)
      }
    }

//    start = System.currentTimeMillis()
    if(node.rightNode.isDefined && !visited.contains(node.rightNode.get)){
      visited.add(node.rightNode.get)
//      println(System.currentTimeMillis()-start)
      // + 逼近上方
      if(value(node.depth%node.k)+lastDist>node.node.vec(node.depth%node.k)){
        searchNode(value, node.rightNode.get, results, visited, K)
      }
    }
//    println(node.depth+" "+(System.currentTimeMillis()-start))
  }

}

case class KdTmpNode(id: Long, dist: Double)  extends Ordered[KdTmpNode] {
  override def compare(that: KdTmpNode): Int = {
    if (dist < that.dist) -1
    else if (dist > that.dist) 1
    else 0
  }
}

object KdTree {

  private[KdTree] def createKdTree(nodes: List[Node], depth: Int, k: Int): Option[KdNode] = {
    if (nodes.size == 0) None
    else {
      // TODO: 只需要查找中位数，后期可以优化
      val sortedNodes = nodes.sortWith { (a, b) => a.vec(depth % k).compareTo(b.vec(depth % k)) < 0}
      val middle = sortedNodes.size / 2
      val left = sortedNodes.slice(0, middle)
      val right = sortedNodes.slice(middle + 1, sortedNodes.size)
      val node = Some(KdNode(sortedNodes(middle), depth, k))
      if (left.size > 0) {
        node.get.leftNode = createKdTree(left, depth + 1, k)
        node.get.leftNode.get.parent = node
      }
      if (right.size > 0) {
        node.get.rightNode = createKdTree(right, depth + 1, k)
        node.get.rightNode.get.parent = node
      }
      node
    }
  }

  private var tree:Option[KdTree] = None

  def createTree(k: Int, file: String) = {
    val nodes = Node.loadFromFile(file)
    println("loading file OK")
    tree = Some(new KdTree(k))
    tree.get.init(nodes)
  }

  // 更高效
  def nearest_1(target: Vector[Double]) = {
    val result = new KdTmpNode(-1l, Double.MaxValue)
    val ans = MutableOpt.from(Some(result))
    val rect = new Rect(target.length)
    (0 until target.length).foreach{
      ind =>
        rect.min.vec(ind) = Double.NegativeInfinity
        rect.max.vec(ind) = Double.PositiveInfinity
    }
    KdNode.nnbr2(tree.get.root, target.toArray, rect, ans, 1)
    ans.get
  }

  def size: Int = size(tree.get.root)

  private def size(root: Option[KdNode]): Int = {
    if(root.isDefined){
      1+size(root.get.leftNode)+size(root.get.rightNode)
    } else {
      0
    }
  }

}

case class KdNode(node: Node, depth: Int, k: Int) {
  var leftNode: Option[KdNode] = None
  var rightNode: Option[KdNode] = None
  var parent: Option[KdNode] = None
}

object KdNode {
  def compareTo(depth: Int, k: Int, value: Vector[Double], treeNode: Vector[Double]) = {
    value(depth % k) - treeNode(depth % k)
  }


  def nnbr(node: Option[KdNode], target: Array[Double], rect: Rect, result: MutableOpt[KdTmpNode], K: Int): Unit ={
    if(node.isDefined) {
      val s = node.get.depth%node.get.k
      val left_rect: Rect = rect
      val right_vect = new Rect(rect)
      left_rect.max.vec(s) = node.get.node.vec(s)
      right_vect.min.vec(s) = node.get.node.vec(s)
      val target_in_left = target(s)<node.get.node.vec(s)

      var nearer_kd: Option[KdNode] = None
      var nearer_hr: Option[Rect] = None
      var future_kd: Option[KdNode] = None
      var future_hr: Option[Rect] = None
      if(target_in_left){
        nearer_kd = node.get.leftNode
        nearer_hr = Some(left_rect)
        future_hr = Some(right_vect)
        future_kd = node.get.rightNode
      } else {
        nearer_kd = node.get.rightNode
        nearer_hr = Some(right_vect)
        future_hr = Some(left_rect)
        future_kd = node.get.leftNode
      }

      nnbr(nearer_kd, target, nearer_hr.get, result, K)

      val p = future_hr.get.closest(target)
      if(Node.euclideanDistance(p.toVector,target.toVector)<result.get.dist){
        val dist_to_target = Node.euclideanDistance(target.toVector, node.get.node.vec)
        if(dist_to_target<result.get.dist){
          result.set(new KdTmpNode(node.get.node.id, dist_to_target))
        }

        nnbr(future_kd,target,future_hr.get, result, K)
      }

    }
  }


  def nnbr2(node: Option[KdNode], target: Array[Double], rect: Rect, result: MutableOpt[KdTmpNode], K: Int): Unit ={
    if(node.isDefined) {
      val s = node.get.depth%node.get.k
      val target_in_left = target(s)<node.get.node.vec(s)
      val old_left = rect.max.vec(s)
      val old_right = rect.min.vec(s)
      if(target_in_left){
        rect.max.vec(s) = node.get.node.vec(s)
        nnbr2(node.get.leftNode, target, rect, result, K)
        rect.max.vec(s) = old_left
        rect.min.vec(s) = node.get.node.vec(s)
        if((target(s)-node.get.node.vec(s)).abs<result.get.dist){
          val dist_to_target = Node.euclideanDistance(target, node.get.node.vec)
          if(dist_to_target<result.get.dist){
            result.set(new KdTmpNode(node.get.node.id, dist_to_target))
          }
          nnbr2(node.get.rightNode,target,rect, result, K)
        }
        rect.min.vec(s) = old_right
      } else {
        rect.min.vec(s) = node.get.node.vec(s)
        nnbr2(node.get.rightNode, target, rect, result, K)
        rect.min.vec(s) = old_right
        rect.max.vec(s) = node.get.node.vec(s)
        //        val p = rect.closest(target)
        //        if(Node.euclideanDistance(p,target)<result.get.dist){
        if((target(s)-node.get.node.vec(s)).abs<result.get.dist){
          val dist_to_target = Node.euclideanDistance(target, node.get.node.vec)
          if(dist_to_target<result.get.dist){
            result.set(new KdTmpNode(node.get.node.id, dist_to_target))
          }
          nnbr2(node.get.leftNode,target,rect, result, K)
        }
        rect.max.vec(s) = old_left
      }
    }
  }

}

case class Node(vec: Vector[Double], id: Long)

object Node {
  def euclideanDistance(a: Vector[Double], b: Vector[Double]): Double = {
    var distance: Double = 0.0
    (0 until a.size).foreach { ind =>
      distance += (a(ind) - b(ind)) * (a(ind) - b(ind))
    }
    // TODO: 其实不需要
//    distance
    Math.sqrt(distance)
  }

  def euclideanDistance(a: Array[Double], b: Array[Double]): Double = {
    var distance: Double = 0.0
    (0 until a.size).foreach { ind =>
      distance += (a(ind) - b(ind)) * (a(ind) - b(ind))
    }
    // TODO: 其实不需要
    //    distance
    Math.sqrt(distance)
  }

  def euclideanDistance(a: Array[Double], b: Vector[Double]): Double = {
    var distance: Double = 0.0
    (0 until a.size).foreach { ind =>
      distance += (a(ind) - b(ind)) * (a(ind) - b(ind))
    }
    // TODO: 其实不需要
    //    distance
    Math.sqrt(distance)
  }

  // ugly but more effective for specific context
  def euclideanDistance(a: Node, b: Node, max: Double) = {
    var distance: Double = 0.0
    val threshold = max * max
    for (ind <- 0 to a.vec.size - 1 if distance < threshold) {
      distance += (a.vec(ind) - b.vec(ind)) * (a.vec(ind) - b.vec(ind))
    }
    if (distance > threshold) {
      Double.NaN
    } else {
      // TODO: 其实不需要
      Math.sqrt(distance)
    }
  }

  def loadFromFile(file: String) = {
    println("loading file")
    var count = 0
    val nodes: ListBuffer[Node] = new ListBuffer[Node]
    for(line <- Source.fromFile(new File(file)).getLines().buffered) {
      nodes +=  nodeParse(line)
      count+=1
      if(count%1000==0){
        println(count)
      }
    }
    nodes.toList
  }

  def nodeParse(line: String): Node = {
    val idVec = line.split(" ")
    val vec = (idVec(1).split(",") map {
      _.toDouble
    }).toVector
    Node(vec, idVec(0).toLong)
  }
}


class MutableOpt[A] {
  private[this] var myValue: A = _
  private[this] var loaded = false
  private def valueEquals(o: Any) = myValue == o
  def get = if (loaded) myValue else throw new NoSuchElementException("MutableOpt")
  def set(a: A): this.type = { loaded = true; myValue = a; this }
  def getOrSet(a: => A): A = {
    if (!loaded) {
      myValue = a
      loaded = true
    }
    myValue
  }
  def isEmpty = !loaded
  def nonEmpty = loaded
  def foreach[U](f: A => U): Unit = if (loaded) f(myValue)
  def transform(f: A => A): this.type = { if (loaded) myValue = f(myValue); this }
  def clear(): this.type = { loaded = false; this }
  def toOption = if (loaded) Some(myValue) else None
  override def toString = if (loaded) "MutableOpt("+myValue.toString+")" else "MutableOpt()"
  override def hashCode = if (loaded) myValue.hashCode else 1751
  override def equals(o: Any) = o match {
    case m: MutableOpt[_] =>
      (isEmpty && m.isEmpty) || (nonEmpty && m.nonEmpty && m.valueEquals(myValue))
    case _ => false
  }
}

object MutableOpt {
  def from[A](o: Option[A]) = {
    val m = new MutableOpt[A]
    o match {
      case Some(a) => m set a
      case _ =>
    }
    m
  }
}
